package repo

import (
	"context"
	"net/http"
	"net/url"
	"os"
	"time"

	"gitee.com/linxing_3/sr-core/yal/yconv"
	"gitee.com/linxing_3/sr-core/yal/ymap"
	"gitee.com/mkwchecking/sr-admin/apps/knowledge/conf"
	"gitee.com/mkwchecking/sr-admin/pkg/file/visit"

	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/documentloaders"
	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/llms/ollama"
	"github.com/tmc/langchaingo/prompts"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/textsplitter"
	"github.com/tmc/langchaingo/vectorstores"
	"github.com/tmc/langchaingo/vectorstores/weaviate"
)

func NewLangchain() (*LangChain, error) {

	llm, err := ollama.New(
		ollama.WithServerURL(conf.Conf.Ollama.Addr),
		ollama.WithModel(conf.Conf.Ollama.EmbedModel),
		ollama.WithHTTPClient(&http.Client{
			Timeout: time.Hour,
		}),
	)
	if err != nil && conf.Conf.Ollama.IgnoreAble {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	// namesapce := uuid.NewString()
	embeder, err := embeddings.NewEmbedder(llm)
	if err != nil {
		return nil, err
	}

	// store, err := chroma.New(
	// 	chroma.WithChromaURL(conf.Conf.VectorDB.Addr),
	// 	chroma.WithDistanceFunction(chroma_go.COSINE),
	// 	chroma.WithNameSpace("llm-embedding-namespace"),
	// 	chroma.WithEmbedder(embeder),
	// )
	url, err := url.Parse(conf.Conf.VectorDB.Addr)
	if err != nil {
		return nil, err
	}

	store, err := weaviate.New(
		weaviate.WithScheme(url.Scheme),
		weaviate.WithHost(url.Host),
		weaviate.WithEmbedder(embeder),
		weaviate.WithIndexName("llm_embed"),
	)

	// qdrantUrl, _ := url.Parse(conf.Conf.VectorDB.Addr)

	// store, err := qdrant.New(
	// 	qdrant.WithURL(*qdrantUrl),
	// 	qdrant.WithEmbedder(embeder),
	// 	qdrant.WithCollectionName(namesapce),
	// )

	if err != nil {
		return nil, err
	}

	return &LangChain{
		store: &store,
	}, nil

}

type LangChain struct {
	store vectorstores.VectorStore

	globalNamespace string
}

func (l *LangChain) EmbeddingFile(ctx context.Context, namespace string, file string, meta map[string]any) error {
	fileStat, err := os.Stat(file)
	if err != nil {
		return err
	}

	data, err := os.Open(file)
	if err != nil {
		return err
	}

	fileType := visit.ParseFileType(file)

	var page documentloaders.Loader
	switch fileType {
	case visit.FileTypeCsv, visit.FileTypeExcel:
		page = documentloaders.NewCSV(data)
	case visit.FileTypePDF:
		page = documentloaders.NewPDF(data, fileStat.Size())
	case visit.FileTypeCode:
		page = documentloaders.NewHTML(data)
	default:
		page = documentloaders.NewText(data)
	}

	split := textsplitter.NewRecursiveCharacter()
	split.ChunkSize = 1000
	split.ChunkOverlap = 30

	docs, err := page.LoadAndSplit(ctx, split)
	if err != nil {
		return err
	}
	for _, doc := range docs {
		yconv.MultiAnyToStr()
		doc.Metadata = ymap.MergeMap(meta, doc.Metadata)
	}

	_, err = l.store.AddDocuments(ctx, docs, vectorstores.WithNameSpace(namespace))
	return err
}

func (l *LangChain) SimilaritySearch(ctx context.Context, query string) ([]schema.Document, error) {
	return l.store.SimilaritySearch(ctx, query, 10)
}

func (l *LangChain) GetRetrieverConversation(ctx context.Context, prompt prompts.FormatPrompter, model string, ns string) (*chains.RetrievalQA, error) {

	retriver := vectorstores.ToRetriever(l.store, 10000, vectorstores.WithNameSpace(ns)) // vectorstores.WithScoreThreshold(0.8),

	llm, err := ollama.New(
		ollama.WithServerURL(conf.Conf.Ollama.Addr),
		ollama.WithModel(model),
	)

	if err != nil {
		return nil, err
	}

	docLLM := chains.NewLLMChain(llm, prompt)
	combinChain := chains.NewStuffDocuments(docLLM)

	qa := chains.NewRetrievalQA(combinChain, retriver)
	return &qa, nil
}
